package cbbx_sm.decision_maker.search;

import java.awt.Dimension;
import java.awt.Point;
import java.awt.geom.Point2D;
import java.text.DecimalFormat;
import java.util.HashSet;

import javax.swing.JFrame;

import cbbx_sm.endtoend.GaussianLargeScale;
import cbbx_sm.utils.ExperimentManager;
import edu.uci.ics.jung.algorithms.layout.StaticLayout;
import edu.uci.ics.jung.graph.DirectedSparseMultigraph;
import edu.uci.ics.jung.graph.util.Pair;
import edu.uci.ics.jung.visualization.VisualizationViewer;
import edu.uci.ics.jung.visualization.control.EditingModalGraphMouse;
import edu.uci.ics.jung.visualization.control.ModalGraphMouse.Mode;
import edu.uci.ics.jung.visualization.decorators.ToStringLabeller;

public class VisualizeSearchGraph {
	
	/**
	 * This class is for the graph visualization.
	 * The visualization toolkit will not accept duplicate edge labels for two different edges.
	 * For this reason this class was added to allow different edges to be labeled with the same probability.
	 * 
	 * @author ronenvaisenberg
	 *
	 */
	public static class EdgeLabel{

		double transitionProbability;
		int id;
		public EdgeLabel(double transitionProbability, int id) {
			this.transitionProbability = transitionProbability;
			this.id = id;
		}
		@Override
		public String toString() {
			return transitionProbability+"";
		}
		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + id;
			return result;
		}
		@Override
		public boolean equals(Object obj) {
			if (this == obj)
				return true;
			if (obj == null)
				return false;
			if (getClass() != obj.getClass())
				return false;
			EdgeLabel other = (EdgeLabel) obj;
			if (id != other.id)
				return false;
			return true;
		}
	}
	
	private static int width = 1400;
	private static  int height = 800;

	/**
	 * @param args
	 * @throws Exception 
	 */
	public static void main(String[] args) throws Exception {
		//WorldStateTransition.setRealWorldModel2States(13,4);
		//WorldStateTransition worldStateTransition = new WorldStateTransition(2);
		ExperimentManager.useNonSeenMotion = true;
		int zoomRegions = 7;
		int levels = 2;
		boolean onlyBestActions = true;
		double upUtility = 0;
		
		WorldStateTransition worldStateTransition = 
			GaussianLargeScale.getGaussianWorldStateTransition(1, zoomRegions);
		
		double[] CiEvent = new double[worldStateTransition.getSchedules().length-1];
		for (int i=0; i<CiEvent.length; i++) {
			CiEvent[i] = EventFactory.getInstance().values()[0];
			//CiEvent[i] = Event.p0;
		}
		//CiEvent[2] = EventFactory.getInstance().values()[7];
		

		
		//WorldStateTransition.setDemoWorldTransitions();
		

		DirectedSparseMultigraph<ScheduleConfigurationNode,EdgeLabel> g = new DirectedSparseMultigraph<ScheduleConfigurationNode,EdgeLabel>();
		
		
		JFrame frame = new JFrame("Test");

		StaticLayout layout = new StaticLayout(g);
		VisualizationViewer<ScheduleConfigurationNode,EdgeLabel> vv = new VisualizationViewer<ScheduleConfigurationNode,EdgeLabel>(layout);

        layout.setSize(new Dimension(width,height));
        
        vv = new VisualizationViewer<ScheduleConfigurationNode,EdgeLabel>(layout);
        vv.getRenderContext().setVertexLabelTransformer(new ToStringLabeller());
        vv.getRenderContext().setEdgeLabelTransformer(new ToStringLabeller());
        vv.setPreferredSize(new Dimension(width,height));
        
        EditingModalGraphMouse gm = new EditingModalGraphMouse(vv.getRenderContext(), 
        		null,
                null);
        
        gm.remove(gm.getPopupEditingPlugin());
        gm.setMode(Mode.PICKING);
        vv.setGraphMouse(gm);
        
        frame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        frame.getContentPane().add(vv);
        frame.pack();
        frame.setVisible(true);
        
        
        addNodes(g, layout, CiEvent, worldStateTransition, levels, onlyBestActions, upUtility);
        
		ScheduleConfigurationNode root = g.getVertices().iterator().next();
		while (root.getParent()!=null){
			root = root.getParent();
		}
		System.out.println("\\documentclass[a4paper,12pt]{article}");
		System.out.println("\\usepackage{pstricks,pst-node,pst-tree}");
		System.out.println("\\begin{document}");
		System.out.println("$");
		System.out.println("\\psset{armB=0cm, levelsep=6cm, treesep=0.15cm, radius=1pt,npos=0.5,nrot=:U}");
        printPSGraph(g, root,1, worldStateTransition);
		System.out.println("$");
		System.out.println("\\end{document}");	
	}

	private static void printPSGraph(
	    DirectedSparseMultigraph<ScheduleConfigurationNode, EdgeLabel> g, 
	    ScheduleConfigurationNode root, int level, WorldStateTransition worldStateTransition) {		
//		if (g.getSuccessorCount(root)==0){
//			if (level==1){
//				System.out.print("\\pstree[treemode=R]{\\Tcircle{"+root.getSelectedAction()+"\\ "+new SchedulerState(root.CiEvent).toString()+"}");
//				DecimalFormat twoDForm = new DecimalFormat("#.####");
//				Double reward = Double.valueOf(twoDForm.format(root.utility+(float)root.getScore()));
//				System.out.print("\\naput{P="+(float)root.transitionProbability+"\\ R="+(reward)+"}}");
//			} else {
//				System.out.print("\\Tc{2pt}");
//			}
//		} else {
			if (level==1){
				System.out.print("\\pstree[treemode=R]{\\Tcircle{"+(root.getSelectedAction()==null?"":root.getSelectedAction()+"\\ ")+"\\ "+/*new SchedulerState(root.CiEvent).toString()+*/"}}");
			} else if (level==2){
				//System.out.print("\\pstree[treemode=R]{\\Tcircle{"+root.getSelectedAction()+"}");
				System.out.print("\\pstree[treemode=R]{\\Tcircle{"+root.getSelectedAction()+"\\ "+/*new SchedulerState(root.CiEvent).toString()+*/"}");

				DecimalFormat twoDForm = new DecimalFormat("#.####");
				Double reward = Double.valueOf(twoDForm.format(root.utility+(float)root.getScore()));
				System.out.print("\\naput{P="+(float)root.transitionProbability+"\\ R="+(reward)+"}}");
				
			}
			else{
				System.out.print("\\pstree[treemode=R]{\\Tc{2pt}}");
			}

			System.out.print("{");

			for (String schedule: worldStateTransition.getSchedules()){
				for (ScheduleConfigurationNode s: g.getSuccessors(root)){
					if (s.getSelectedAction().compareTo(schedule)==0){
						printPSGraph(g, s,level+1, worldStateTransition);
					}
				}
			}

			System.out.print("}");
			
		//}	
	}

	private static void addNodes(DirectedSparseMultigraph<ScheduleConfigurationNode, EdgeLabel> g, StaticLayout layout,double[] CiEvent, 
	    WorldStateTransition worldStateTransition, int levels, boolean onlyBestActions, double beta) {
		
		double prob = 0;
		
		ScheduleConfigurationNode root = 
			new ScheduleConfigurationNode(null, CiEvent, null,worldStateTransition, prob,1,beta);
		TreeLookAhead.growTree(root,worldStateTransition, levels);
		
        // Add Root.
		g.addVertex(root);
        Point2D p = new Point(0,height/2);
		layout.setLocation(root, p);
		
		SearchDriver driver = new SearchDriver(root);
		
		HashSet<SearchNode> frontierNodes = null;
		int id = 0;
		for (int i=1; i<=levels; i++){
			frontierNodes = driver.expandFrontierNodes(onlyBestActions);
			int j=1;
			

			
			for (SearchNode n: frontierNodes){
		        // Add Node.
				// Reason behind checking if the graph does not contain it is to avoid nodes which are added to the frontier due to pruning of their path.
				if (!g.containsVertex((ScheduleConfigurationNode)n)){
					g.addVertex((ScheduleConfigurationNode)n);
			        p = new Point(i*((width-300)/levels),0+(j++)*Math.max(5,((height-30)/frontierNodes.size())));
					layout.setLocation(n, p);
					
					ScheduleConfigurationNode parent = ((ScheduleConfigurationNode) n).getParent();
					
					// Add edge.
					EdgeLabel e1 = new EdgeLabel(worldStateTransition.getTransitionProbability(parent.CiEvent, ((ScheduleConfigurationNode) n).CiEvent,n.getSelectedAction()),id++);
					g.addEdge(e1, new Pair(parent,n));	
				}
			}
		}		
	}

}
